from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnableConfig

from base import model

# 聊天机器人案例
# 机器人记住聊天记录
# 流式输出 防止内存溢出，避免一次性将所有消息加载到内存中。可以使用生成器来实现这个功能：

# 全局 消息key
msgKey = "msg_key"

# 初始化大模型

# 聊天提示模板
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个{lang}智能开发助手"),
        MessagesPlaceholder(
            variable_name=msgKey
        )
    ]
)


# 执行链
chain = prompt | model
# 用于保存每个用户session对应的聊天历史
# key->sessionId
# value -> 历史聊天对象
historyStore = {}

# 获取对应session的聊天历史
def get_session_history(session_id: str):
    # 如果不存在 则创建一个
    if session_id not in historyStore:
        # history = ConversationBufferMemory()
        history = ChatMessageHistory()
        historyStore[session_id] = history
        return history
    return historyStore.get(session_id)


doMessage = RunnableWithMessageHistory(
    chain,
    get_session_history=get_session_history,
    # 指定当前消息的key
    # 每次聊天时发送这个key
    input_messages_key=msgKey
)

# 开始聊天
config = RunnableConfig(
    configurable={'session_id': 'zhangsan'},
)
# 第一轮

result = doMessage.invoke(
    {
        msgKey: [HumanMessage(content="帮我写一个HelloWorld程序")],
        'lang': 'Java'
    },
    config
)
print(result)

# 第二轮
result = doMessage.invoke(
    {
        msgKey: [HumanMessage(content="加上注释")],
        'lang': 'Java'
    },
    config
)
print(result)
print('------------------------------------------------------------------')
#第三轮 返回流式输出
streamResult = doMessage.stream(
    {
        msgKey: [HumanMessage(content="类上 加上'@author 张治保'的注释")],
        'lang': 'Java'
    },
    config
)

for token in streamResult:
    #每次响应都是一个token
    print(token,end='_')
# 第四论
